import numpy
import random
#T=2Na/3R*K

def kinetic(m, v):
  return numpy.sum(m * (v * v).sum(axis=1) * 0.5)
def temperature(m, v):
  Na = 6.022e23
  R=8.314
  k = 1.38064852e-23
  print(m.shape, v.shape)
  kstd = kinetic(m, v) / (Na * 1000) * 1e10
  #print(kstd)
  return kstd / k / 7051 * 2 / 3
import pandas
atoms = pandas.read_csv('project.dump', delimiter=' ')
v = atoms[['vx', 'vy', 'vz']].values
m = atoms['mass'].values
#v is in A/fs, m is in grams/mole
#v: 1e-10m/1e-15s = 1e5m/s
#m: m / Na / 1000
# vstd = v * 1e5
# mstd = m / Na / 1000
def distnormal(mu, sigma):
  u = 1 - numpy.random.rand()
  v = 1 - numpy.random.rand()
  r = numpy.sqrt(-2 * numpy.log(u))
  theta = 2 * numpy.pi * v
  z = r * numpy.sin(theta)
  return z * numpy.sqrt(sigma) + mu
def create_velocity(m, t):
  Na = 6.022e23
  R=8.314
  k = 1.38064852e-23
  v = numpy.zeros((m.shape[0], 3))
  for i in range(m.shape[0]):
    v[i][0] = distnormal(0, k * t * Na * 1000 / m[i]) / 1e5
    v[i][1] = distnormal(0, k * t * Na * 1000 / m[i]) / 1e5
    v[i][2] = distnormal(0, k * t * Na * 1000 / m[i]) / 1e5
  t0 = temperature(m, v)
  scale = numpy.sqrt(t / t0)
  v = v * scale
  return v